import tensorflow as tf
import matplotlib.pylab as plt
import numpy as np
from tqdm import tqdm
import pickle
import random
from sklearn.utils import shuffle
import cv2

numElems = 100
data_len = 10000

mnist = tf.keras.datasets.mnist
(x_train_original, y_train_original), (_, _) = mnist.load_data()
x_train = []
y_train = []
x_test = []
y_test = []
# 截断，方便调试
x_train_original = x_train_original[:data_len]
y_train_original = y_train_original[:data_len]
print("原始x", x_train_original.shape)
print("原始y", y_train_original.shape)

for i in tqdm(range(data_len)):
    image = x_train_original[i]
    target = y_train_original[i]

    "数据预处理"
    nonzero_array = np.nonzero(image)
    nonzero_count = len(nonzero_array[0])
    if nonzero_count > numElems:  # 如果目标大于100维，随机取100维
        idx = np.round(np.linspace(0, nonzero_count - 1, numElems)).astype(int)  # 随机均匀获得100维坐标
        image[image >= 0] = 0
        for each_index in idx:
            x = nonzero_array[0][each_index]
            y = nonzero_array[1][each_index]
            image[x][y] = 1
    else:  # 如果目标小于100维，随机不全不足的维数
        '图像膨胀'
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        while len(np.nonzero(image)[0]) < 100:
            image = cv2.dilate(image, kernel)  # 膨胀图像
        nonzero_array = np.nonzero(image)
        nonzero_count = len(nonzero_array[0])
        if nonzero_count > numElems:  # 如果目标大于100维，随机取100维
            idx = np.round(np.linspace(0, nonzero_count - 1, numElems)).astype(int)  # 随机均匀获得100维坐标
            image[image >= 0] = 0
            for each_index in idx:
                x = nonzero_array[0][each_index]
                y = nonzero_array[1][each_index]
                image[x][y] = 1

    # 左上角
    image_left_top = np.zeros([56, 56])
    for j in range(0, 28):
        for k in range(0, 28):
            image_left_top[j][k] = image[j][k]
    # 右上角
    image_right_top = np.zeros([56, 56])
    for j in range(0, 28):
        for k in range(28, 56):
            image_right_top[j][k] = image[j][k - 28]
    # 左下角
    image_left_bottom = np.zeros([56, 56])
    for j in range(28, 56):
        for k in range(0, 28):
            image_left_bottom[j][k] = image[j - 28][k]
    # 右下角
    image_right_bottom = np.zeros([56, 56])
    for j in range(28, 56):
        for k in range(28, 56):
            image_right_bottom[j][k] = image[j - 28][k - 28]
    # 压入训练集和测试集
    x_train.append(image_left_top)
    x_test.extend([image_right_top, image_left_bottom, image_right_bottom])
    y_train.append(target)
    y_test.extend([target, target, target])
    if i == 0:
        plt.figure(figsize=(8, 8))
        ax1 = plt.subplot(221)
        plt.imshow(image_left_top)
        plt.title("TrainDataset")
        ax2 = plt.subplot(222)
        plt.imshow(image_right_top)
        plt.title("TestDataset")
        ax3 = plt.subplot(223)
        plt.imshow(image_left_bottom)
        plt.title("TestDataset")
        ax4 = plt.subplot(224)
        plt.imshow(image_right_bottom)
        plt.title("TestDataset")
        plt.savefig("训练集.png")
        plt.show()
# 打印shape
x_train = np.array(x_train, dtype=np.int8)
x_test = np.array(x_test, dtype=np.int8)
y_train = np.array(y_train, dtype=np.int8)
y_test = np.array(y_test, dtype=np.int8)
print("训练x", x_train.shape)
print("训练y", y_train.shape)
print("测试x", x_test.shape)
print("测试y", y_test.shape)
# 打乱
x_train, y_train = shuffle(x_train, y_train, random_state=0)
x_test, y_test = shuffle(x_test, y_test, random_state=0)
data_set = {"x_train": x_train,
            "x_test": x_test,
            "y_train": y_train,
            "y_test": y_test}
with open("data_set.pickle", 'wb') as f:
    pickle.dump(data_set, f)
